package nl.runnable.solr.xslt;

import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.util.Collections;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamReader;
import javax.xml.transform.Source;
import javax.xml.transform.Templates;
import javax.xml.transform.TransformerConfigurationException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.stream.StreamSource;

import nl.runnable.solr.util.RequestHandlerHelper;

import org.apache.solr.common.ResourceLoader;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.util.ContentStream;
import org.apache.solr.request.SolrQueryRequest;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Default {@link TemplatesResolver} implementation.
 * 
 * @author Laurens Fridael
 * 
 */
public class DefaultTemplatesResolver implements TemplatesResolver {

	private final Logger logger = LoggerFactory.getLogger(getClass());

	/* State */

	private Map<QualifiedName, String> stylesheetsByRootElementName = Collections.emptyMap();

	private final Map<String, Templates> templatesByStylesheet = new ConcurrentHashMap<String, Templates>();

	/* Dependencies */

	private RequestHandlerHelper requestHandlerHelper = new RequestHandlerHelper();

	private TransformerFactory transformerFactory;

	private XMLInputFactory xmlInputFactory;

	/* Configuration */

	private String stylesheetPrefix = "";

	private String stylesheetSuffix = "";

	private String defaultStylesheet;

	private boolean disableStylesheetCache = false;
	
	private boolean disableRootElementScan = false;

	/**
	 * Resolves {@link Templates} against the given {@link SolrQueryRequest} and {@link ContentStream}.
	 * <p>
	 * This implementation attempts to resolve the XSL stylesheet using the following steps, in this particular order:
	 * <ol>
	 * <li>Use the stylesheet name specified by the {@link RequestParameters#STYLESHEET} request parameter.</li>
	 * <li>Scan for the XML root element and look up the stylesheet by the root element name.</li>
	 * <li>Use the default stylesheet if one is configured.</li>
	 * </ol>
	 * If none of the above yields a stylesheet name, the implementation returns null.
	 * 
	 * @param request
	 * @param contentStream
	 * @return The matching {@link Templates} or null if none could be determined.
	 */
	@Override
	public Templates resolveTemplates(final SolrQueryRequest request, final ContentStream contentStream)
			throws IOException {
		if (request == null) {
			throw new IllegalArgumentException("SolrRequest cannot be null.");
		}

		final SolrParams params = request.getParams();
		String stylesheet = params.get(RequestParameters.STYLESHEET);
		if (stylesheet == null && getStylesheetsByRootElementName().isEmpty() == false && isDisableRootElementScan() == false) {
			QualifiedName rootElementName = null;
			if (rootElementName == null) {
				if (contentStream == null) {
					throw new IllegalStateException(
							"No ContentStream provided. Cannot scan content for XML root element.");
				}
				rootElementName = scanForRootElementName(contentStream);
			}
			if (rootElementName != null) {
				stylesheet = getStylesheetsByRootElementName().get(rootElementName);
			}
		}
		if (stylesheet == null) {
			stylesheet = getDefaultStylesheet();
		}
		Templates templates = null;
		if (stylesheet != null) {
			stylesheet = String.format("%s%s%s", getStylesheetPrefix(), stylesheet, getStylesheetSuffix());
			templates = getTemplates(request, stylesheet);
		}
		return templates;
	}

	protected Templates getTemplates(final SolrQueryRequest request, final String stylesheet) throws IOException {
		Templates templates = null;
		if (templatesByStylesheet.containsKey(stylesheet) && isDisableStylesheetCache() == false) {
			templates = templatesByStylesheet.get(stylesheet);
		} else {
			templates = loadTemplates(request, stylesheet);
			if (templates != null) {
				templatesByStylesheet.put(stylesheet, templates);
			}
		}
		return templates;
	}

	protected Templates loadTemplates(final SolrQueryRequest request, final String stylesheet) throws IOException {
		Templates templates = null;
		try {
			String systemId = null;
			final URI stylesheetUri = getRequestHandlerHelper().getResource(request, stylesheet);
			final InputStream stylesheetInputStream;
			if (stylesheetUri != null) {
				systemId = stylesheetUri.toString();
				stylesheetInputStream = stylesheetUri.toURL().openStream();
			} else {
				final ResourceLoader resourceLoader = getRequestHandlerHelper().getResourceLoader(request);
				stylesheetInputStream = resourceLoader.openResource(stylesheet);
			}
			final Source source = new StreamSource(stylesheetInputStream, systemId);
			templates = getTransformerFactory().newTemplates(source);
		} catch (final TransformerConfigurationException e) {
			logger.warn("Error reading stylesheet {}.", new Object[] { stylesheet, e });
		} catch (final RuntimeException e) {
			/* Thrown by ResourceLoader.openStream(). */
			logger.warn("Error reading stylesheet {}.", new Object[] { stylesheet, e });
		}
		return templates;
	}

	/**
	 * Scans the given {@link ContentStream} for the root element name.
	 * 
	 * @param contentStream
	 * @return The {@link QualifiedName} for the root element.
	 */
	protected QualifiedName scanForRootElementName(final ContentStream contentStream) throws IOException {
		XMLStreamReader xmlStreamReader = null;
		try {
			xmlStreamReader = getXmlInputFactory().createXMLStreamReader(contentStream.getStream());
			QualifiedName rootElementName = null;
			if (xmlStreamReader.nextTag() == XMLStreamReader.START_ELEMENT) {
				rootElementName = new QualifiedName(xmlStreamReader.getLocalName(), xmlStreamReader.getNamespaceURI());
			}
			return rootElementName;
		} catch (final XMLStreamException e) {
			throw new RuntimeException(e);
		} finally {
			if (xmlStreamReader != null) {
				try {
					xmlStreamReader.close();
				} catch (final XMLStreamException e) {
					throw new RuntimeException(e);
				}
			}
		}
	}

	/* Dependencies */

	public void setRequestHandlerHelper(final RequestHandlerHelper requestHandlerHelper) {
		this.requestHandlerHelper = requestHandlerHelper;
	}

	protected RequestHandlerHelper getRequestHandlerHelper() {
		return requestHandlerHelper;
	}

	protected XMLInputFactory getXmlInputFactory() {
		if (xmlInputFactory == null) {
			xmlInputFactory = XMLInputFactory.newInstance();
			xmlInputFactory.setProperty(XMLInputFactory.IS_NAMESPACE_AWARE, true);
		}
		return xmlInputFactory;
	}

	protected TransformerFactory getTransformerFactory() {
		if (transformerFactory == null) {
			transformerFactory = TransformerFactory.newInstance();
		}
		return transformerFactory;
	}

	/* Configuration */

	public void setStylesheetsByRootElementName(final Map<QualifiedName, String> stylesheetsByRootElementName) {
		this.stylesheetsByRootElementName = stylesheetsByRootElementName;
	}

	protected Map<QualifiedName, String> getStylesheetsByRootElementName() {
		return stylesheetsByRootElementName;
	}

	public void setStylesheetPrefix(final String stylesheetPrefix) {
		if (stylesheetPrefix == null) {
			throw new IllegalArgumentException("Stylesheet prefix cannot be null.");
		}
		this.stylesheetPrefix = stylesheetPrefix;
	}

	protected String getStylesheetPrefix() {
		return stylesheetPrefix;
	}

	public void setStylesheetSuffix(final String stylesheetSuffix) {
		if (stylesheetSuffix == null) {
			throw new IllegalArgumentException("Stylesheet suffix cannot be null.");
		}
		this.stylesheetSuffix = stylesheetSuffix;
	}

	protected String getStylesheetSuffix() {
		return stylesheetSuffix;
	}

	public void setDefaultStylesheet(final String defaultStylesheet) {
		this.defaultStylesheet = defaultStylesheet;
	}

	protected String getDefaultStylesheet() {
		return defaultStylesheet;
	}

	public void setDisableStylesheetCache(final boolean disableStylesheetCache) {
		this.disableStylesheetCache = disableStylesheetCache;
	}

	public boolean isDisableStylesheetCache() {
		return disableStylesheetCache;
	}
	
	public void setDisableRootElementScan(boolean disableRootElementScan) {
	  this.disableRootElementScan = disableRootElementScan;
	}
	
	public boolean isDisableRootElementScan() {
	  return disableRootElementScan;
	}

}
